import torch
import numpy as np
import viser
import time
from scipy.spatial.transform import Rotation

class PointCloudViewer:
    """A Viser-based interactive viewer for inspecting multiple registered point clouds and cameras."""
    def __init__(self, pc_list, color_list, conf_list, cam_dict, view_names=None, port=8080, show_camera=True, point_size=0.002):
        self.server = viser.ViserServer(port=port)
        # --- FIX: Set the world's "up" direction to match the dataset's coordinate system (Y-down) ---
        self.server.set_up_direction("-y")

        self.pcs, self.all_steps = self._read_data(pc_list, color_list, conf_list)
        self.cam_dict = cam_dict
        self.view_names = view_names if view_names is not None else [f"View {i}" for i in range(len(pc_list))]
        self.num_frames = len(self.all_steps)
        self.show_camera = show_camera
        
        self.psize_slider = self.server.gui.add_slider("Point Size", min=0.0001, max=0.01, step=0.0001, initial_value=point_size)
        self.camsize_slider = self.server.gui.add_slider("Camera Size", min=0.01, max=0.5, step=0.01, initial_value=0.05)
        
        # Modified: Add reset button and dropdown with on_update event for dropdown
        with self.server.add_gui_folder("Reset View"):
            self.reset_button = self.server.gui.add_button("Reset to View:")
            self.view_dropdown = self.server.gui.add_dropdown("Select View", options=self.view_names, initial_value=self.view_names[0])

            @self.reset_button.on_click
            def _(event: viser.GuiEvent) -> None:
                """Resets the camera to a viewpoint calculated to best see the selected view's points."""
                assert event.client is not None
                selected_view = self.view_dropdown.value
                view_index = self.view_names.index(selected_view)
                camera_pos, camera_target = self._calculate_viewpoint_for_view(view_index)
                event.client.camera.position = camera_pos
                event.client.camera.look_at = camera_target

            @self.view_dropdown.on_update
            def _(event: viser.GuiEvent) -> None:
                """Updates the camera to a viewpoint calculated to best see the selected view's points when dropdown changes."""
                assert event.client is not None
                selected_view = self.view_dropdown.value
                view_index = self.view_names.index(selected_view)
                camera_pos, camera_target = self._calculate_viewpoint_for_view(view_index)
                event.client.camera.position = camera_pos
                event.client.camera.look_at = camera_target
            
        self.pc_handles = []
        self.cam_handles = []
        self.frame_nodes = []
        self.visibility_checkboxes = []

        @self.psize_slider.on_update
        def _(slider: viser.GuiInputHandle) -> None:
            for handle in self.pc_handles:
                handle.point_size = slider.value

        @self.camsize_slider.on_update
        def _(slider: viser.GuiInputHandle) -> None:
            for handle in self.cam_handles:
                handle.scale = slider.value
        
        self.initial_camera_pos, self.initial_camera_target = self._calculate_initial_viewpoint()

        self.server.on_client_connect(self._on_client_connect)

    def _calculate_initial_viewpoint(self):
        """Calculates a good initial viewpoint to see the first two point clouds."""
        visible_points = []
        for i in range(min(2, self.num_frames)):
            data = self.pcs[i]
            pts, _ = self._parse_pc_data(data["pc"], data["color"], data["conf"])
            visible_points.append(pts)
        
        if not visible_points:
            return np.array([0.0, 0.0, -1.0]), np.array([0.0, 0.0, 0.0])

        all_points = np.concatenate(visible_points, axis=0)
        if all_points.shape[0] == 0:
            return np.array([0.0, 0.0, -1.0]), np.array([0.0, 0.0, 0.0])

        center = np.mean(all_points, axis=0)
        
        # Calculate a camera position offset from the center
        # We'll move back along the Z-axis and up along the Y-axis (which is Y-down)
        # The distance will be based on the spread of the points
        max_dist = np.max(np.linalg.norm(all_points - center, axis=1))
        
        # Heuristic for camera position
        camera_pos = center + np.array([0.0, -max_dist * 0.5, -max_dist * 2.0])
        
        return camera_pos, center

    def _calculate_viewpoint_for_view(self, view_index):
        """Calculates a viewpoint to best see all points of a specific view."""
        if view_index >= self.num_frames:
            return np.array([0.0, 0.0, -1.0]), np.array([0.0, 0.0, 0.0])

        data = self.pcs[view_index]
        pts, _ = self._parse_pc_data(data["pc"], data["color"], data["conf"])
        
        if pts.shape[0] == 0:
            return np.array([0.0, 0.0, -1.0]), np.array([0.0, 0.0, 0.0])

        center = np.mean(pts, axis=0)
        max_dist = np.max(np.linalg.norm(pts - center, axis=1))
        camera_pos = center + np.array([0.0, -max_dist * 0.5, -max_dist * 2.0])
        
        return camera_pos, center

    def _on_client_connect(self, client: viser.ClientHandle) -> None:
        """Callback for when a new client connects. Sets the initial camera position."""
        if self.num_frames > 0:
            client.camera.position = self.initial_camera_pos
            client.camera.look_at = self.initial_camera_target

    def set_camera_view(self, client: viser.ClientHandle, view_index: int) -> None:
        """Sets the client's camera to the pose of a specific view."""
        if view_index >= self.num_frames:
            return

        R = self.cam_dict["R"][view_index]
        t = self.cam_dict["t"][view_index]

        # --- FIX: Use the raw CV pose directly, without any GL conversion ---
        pose = np.eye(4)
        pose[:3, :3], pose[:3, 3] = R, t

        rotation_matrix = pose[:3, :3]
        translation_vector = pose[:3, 3]

        # Pull the camera slightly back from its origin to ensure points are visible
        # The camera's forward vector is its local Z-axis (third column of rotation matrix)
        forward_vector = rotation_matrix[:, 2]
        # Move backward, which is the opposite of the forward direction
        new_position = translation_vector - forward_vector * 0.1 

        r = Rotation.from_matrix(rotation_matrix)
        quat_xyzw = r.as_quat()
        quat_wxyz = np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]])

        client.camera.wxyz = quat_wxyz
        client.camera.position = new_position

    def _read_data(self, pc_list, color_list, conf_list):
        pcs = {}
        step_list = []
        for i, pc in enumerate(pc_list):
            pcs[i] = {"pc": pc, "color": color_list[i], "conf": conf_list[i]}
            step_list.append(i)
        return pcs, step_list

    def _parse_pc_data(self, pc, color, conf):
        pred_pts, color, conf = pc.reshape(-1, 3), color.reshape(-1, 3), conf.reshape(-1)
        mask = conf > 0.5
        return pred_pts[mask], color[mask]

    def _add_pc(self, step):
        data = self.pcs[step]
        pred_pts, color = self._parse_pc_data(data["pc"], data["color"], data["conf"])
        self.pc_handles.append(self.server.add_point_cloud(
            name=f"/frames/{step}/pred_pts", points=pred_pts, colors=color, point_size=self.psize_slider.value
        ))

    def _add_camera(self, step):
        cam = self.cam_dict
        focal, (pp_x, pp_y), R, t = cam["focal"][step], cam["pp"][step], cam["R"][step], cam["t"][step]
        
        # --- FIX: The frustum object itself still needs a transform to display correctly ---
        # A camera frustum in viser is an object that expects an OpenGL-style pose.
        pose = np.eye(4)
        pose[:3, :3], pose[:3, 3] = R, t
        cv_to_opengl = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
        pose_gl = pose @ cv_to_opengl

        fov = 2 * np.arctan(pp_y / focal)
        aspect = pp_x / pp_y

        r = Rotation.from_matrix(pose_gl[:3, :3])
        quat_xyzw = r.as_quat()
        quat_wxyz = np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]])
        
        self.cam_handles.append(self.server.add_camera_frustum(
            name=f"/frames/{step}/camera", fov=fov, aspect=aspect, wxyz=quat_wxyz,
            position=pose_gl[:3, 3], scale=self.camsize_slider.value, color=(230, 120, 50)
        ))

    def run(self):
        for i in range(self.num_frames):
            frame_node = self.server.add_frame(f"/frames/{i}", show_axes=False)
            self.frame_nodes.append(frame_node)
            self._add_pc(i)
            if self.show_camera:
                self._add_camera(i)
            frame_node.visible = i < 2

        with self.server.add_gui_folder("View Visibility"):
            all_on = self.server.gui.add_button("All On")
            all_off = self.server.gui.add_button("All Off")

            @all_on.on_click
            def _(event: viser.GuiEvent) -> None:
                for cb in self.visibility_checkboxes: cb.value = True
            
            @all_off.on_click
            def _(event: viser.GuiEvent) -> None:
                for cb in self.visibility_checkboxes: cb.value = False

            for i in range(self.num_frames):
                checkbox = self.server.gui.add_checkbox(f"{self.view_names[i]}", initial_value=i < 2)
                self.visibility_checkboxes.append(checkbox)

                def make_handler(index, cb_handle):
                    def handler(event: viser.GuiEvent) -> None:
                        self.frame_nodes[index].visible = cb_handle.value
                    return handler
                checkbox.on_update(make_handler(i, checkbox))
        while True:
            time.sleep(1.0)

def launch_visualizer(pc_list, color_list, conf_list, cam_dict, view_names=None, port=8080, show_camera=False):
    viewer = PointCloudViewer(
        pc_list=pc_list, color_list=color_list, conf_list=conf_list,
        cam_dict=cam_dict, view_names=view_names, port=port, show_camera=show_camera
    )
    viewer.run()